#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include "cqupt.h"
#pragma comment(lib, "Ws2_32.lib")

#define CQUPT_EXAMPLE_SERVER_IP "127.0.0.1"


int cqupt_example_upload_file(int sockfd, char* file_name) {
    if (file_name == NULL || strlen(file_name) == 0) {
        fprintf(stderr, "Invalid file name.\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }
    
    size_t filename_len = strlen(file_name);
    if (filename_len > CQUPT_MAX_PAYLOAD_SIZE) {
        fprintf(stderr, "File name too long.\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }
    
    // 第一个消息：发送文件名
    cqupt_message_t message = {0};
    message.header.version = CQUPT_VERSION;
    message.header.type = CQUPT_CMD_UPLOAD;
    message.header.length = filename_len;
    message.header.checksum = cqupt_calculate_checksum(&message.header, sizeof(message.header) - sizeof(message.header.checksum));
    memcpy(message.payload, file_name, filename_len);

    if (send(sockfd, (char *)&message, CQUPT_HEADER_SIZE + filename_len, 0) <= 0) {
        fprintf(stderr, "Failed to send file name.\n");
        return -1;
    }
    
    // 读取整个文件内容
    FILE *file = fopen(file_name, "rb");
    if (!file) {
        fprintf(stderr, "Failed to open file: %s\n", file_name);
        return -1;
    }
    
    // 获取文件大小
    fseek(file, 0, SEEK_END);
    long file_size = ftell(file);
    fseek(file, 0, SEEK_SET);
    
    if (file_size > CQUPT_MAX_PAYLOAD_SIZE) {
        fprintf(stderr, "File too large (max %d bytes).\n", CQUPT_MAX_PAYLOAD_SIZE);
        fclose(file);
        return -1;
    }
    
    // 第二个消息：发送文件内容
    memset(&message, 0, sizeof(message));
    message.header.version = CQUPT_VERSION;
    message.header.type = CQUPT_CMD_UPLOAD;
    message.header.length = file_size;
    message.header.checksum = cqupt_calculate_checksum(&message.header, sizeof(message.header) - sizeof(message.header.checksum));
    
    if (fread(message.payload, 1, file_size, file) != file_size) {
        fprintf(stderr, "Failed to read file content.\n");
        fclose(file);
        return -1;
    }
    fclose(file);
    
    if (send(sockfd, (char *)&message, CQUPT_HEADER_SIZE + file_size, 0) <= 0) {
        fprintf(stderr, "Failed to send file content.\n");
        return -1;
    }
    
    printf("File '%s' uploaded successfully.\n", file_name);
    return 0;
}

int cqupt_example_download_file(int sockfd, char* file_name) {
    if (file_name == NULL || strlen(file_name) == 0) {
        fprintf(stderr, "Invalid file name.\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }
    
    size_t filename_len = strlen(file_name);
    if (filename_len > CQUPT_MAX_PAYLOAD_SIZE) {
        fprintf(stderr, "File name too long.\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }
    
    // 发送下载请求消息（文件名）
    cqupt_message_t message = {0};
    message.header.version = CQUPT_VERSION;
    message.header.type = CQUPT_CMD_DOWNLOAD;
    message.header.length = filename_len;
    message.header.checksum = cqupt_calculate_checksum(&message.header, sizeof(message.header) - sizeof(message.header.checksum));
    memcpy(message.payload, file_name, filename_len);
    
    if (send(sockfd, (char *)&message, CQUPT_HEADER_SIZE + filename_len, 0) <= 0) {
        fprintf(stderr, "Failed to send download request.\n");
        return -1;
    }

    // 接收第一个消息：文件名确认
    cqupt_message_t response = {0};
    if (recv(sockfd, (char*)&response.header, CQUPT_HEADER_SIZE, 0) <= 0) {
        fprintf(stderr, "Failed to receive response header.\n");
        return -1;
    }
    
    if (cqupt_validate_header(&response.header) != 0) {
        fprintf(stderr, "Invalid header received.\n");
        return CQUPT_ERROR_INVALID_HEADER;
    }
    
    if (response.header.type == CQUPT_ACK_FAILURE) {
        fprintf(stderr, "File not found on server.\n");
        return -1;
    }
    
    if (response.header.length > 0) {
        if (recv(sockfd, (char*)response.payload, response.header.length, 0) <= 0) {
            fprintf(stderr, "Failed to receive file name confirmation.\n");
            return -1;
        }
        response.payload[response.header.length] = '\0';
        printf("Server confirmed file: %s\n", response.payload);
    }
    
    // 接收第二个消息：文件内容
    if (recv(sockfd, (char*)&response.header, CQUPT_HEADER_SIZE, 0) <= 0) {
        fprintf(stderr, "Failed to receive file content header.\n");
        return -1;
    }
    
    if (cqupt_validate_header(&response.header) != 0) {
        fprintf(stderr, "Invalid file content header received.\n");
        return CQUPT_ERROR_INVALID_HEADER;
    }
    
    if (response.header.length > 0) {
        if (recv(sockfd, (char*)response.payload, response.header.length, 0) <= 0) {
            fprintf(stderr, "Failed to receive file content.\n");
            return -1;
        }
        
        // 写入本地文件
        FILE *file = fopen(file_name, "wb");
        if (!file) {
            fprintf(stderr, "Failed to open file for writing: %s\n", file_name);
            return -1;
        }
        
        if (fwrite(response.payload, 1, response.header.length, file) != response.header.length) {
            fprintf(stderr, "Failed to write file data.\n");
            fclose(file);
            return -1;
        }
        fclose(file);
    }
    
    printf("File '%s' downloaded successfully.\n", file_name);
    return 0;
}

int main(int argc , char *argv[]) {
    if (argc==1 || argc > 3) {
        fprintf(stderr, "Invalid params, use 'help' instead\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }else if (strcmp(argv[1], "help") == 0) {
        printf("Usage: %s <command> <file_name>\n", argv[0]);
        printf("Commands:\n");
        printf("  upload    <file_name>     Upload a file to the server.\n");
        printf("  download  <file_name>     Download a file from the server.\n");
        return 0;
    }else if (strcmp(argv[1], "upload") != 0 && strcmp(argv[1], "download") != 0) {
        fprintf(stderr, "Invalid params, use 'help' instead\n");
        return CQUPT_ERROR_INVALID_PARAM;
    }

    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0) {
        fprintf(stderr, "WSAStartup failed.\n");
        return -1;
    }
    SOCKET client_socket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
    if (client_socket == INVALID_SOCKET) {
        fprintf(stderr, "Socket creation failed.\n");
        WSACleanup();
        return -1;
    }
    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(CQUPT_PORT);
    server_addr.sin_addr.s_addr = inet_addr(CQUPT_EXAMPLE_SERVER_IP);
    if (connect(client_socket, (struct sockaddr *)&server_addr, sizeof(server_addr)) != 0) {
        fprintf(stderr, "Connection to server failed.\n");
        closesocket(client_socket);
        WSACleanup();
        return -1;
    }
    printf("Connected to server at %s:%d\n", CQUPT_EXAMPLE_SERVER_IP, CQUPT_PORT);

    if (strcmp(argv[1], "download") == 0) {
        if (cqupt_example_download_file(client_socket, argv[2]) != 0) {
            fprintf(stderr, "Download failed.\n");
            closesocket(client_socket);
            WSACleanup();
            return -1;
        }
    } else if (strcmp(argv[1], "upload") == 0) {
        if (cqupt_example_upload_file(client_socket, argv[2]) != 0) {
            fprintf(stderr, "Upload failed.\n");
            closesocket(client_socket);
            WSACleanup();
            return -1;
        }
    } else {
        fprintf(stderr, "Invalid command. Use 'upload' or 'download'.\n");
        closesocket(client_socket);
        WSACleanup();
        return CQUPT_ERROR_INVALID_PARAM;
    }
    closesocket(client_socket);
    WSACleanup();
    return 0;
}